package networkTraining.epoch;

import java.util.HashMap;
import java.util.Map;

import networkTraining.IDataSet;
import networkTraining.INetworkTrainer;
import networkTraining.algorithms.ITrainingAlgorithm;

public abstract class EpochExecutorBase implements IEpochExecutor
{
	private INetworkTrainer networkTrainer;

	private int epochIndex;

	private Map<String, Double> totalErrorMap;
	private Map<String, Double> averageErrorMap;
	private Map<String, Double> squaredErrorSumMap;
	private Map<String, Double> rmsErrorMap;
	private double epochRMS;

	public EpochExecutorBase()
	{
		totalErrorMap = new HashMap<String, Double>();
		averageErrorMap = new HashMap<String, Double>();
		squaredErrorSumMap = new HashMap<String, Double>();
		rmsErrorMap = new HashMap<String, Double>();
	}

	public INetworkTrainer getNetworkTrainer()
	{
		return networkTrainer;
	}

	public void setNetworkTrainer(INetworkTrainer networkTrainer)
	{
		this.networkTrainer = networkTrainer;
	}

	@Override
	public void applyTrainingAlgorithm(Map<String, Double> errorMap)
	{
		ITrainingAlgorithm trainingAlgorithm = networkTrainer.getTrainingAlgorithm();
		trainingAlgorithm.trainNeuralNetwork(errorMap);
	}

	@Override
	public int getEpochIndex()
	{
		return epochIndex;
	}

	@Override
	public void incrementEpochIndex()
	{
		epochIndex++;
	}

	@Override
	public Map<String, Double> getTotalErrorMap()
	{
		return totalErrorMap;
	}

	@Override
	public Map<String, Double> getAverageErrorMap()
	{
		return averageErrorMap;
	}

	@Override
	public Map<String, Double> getSquaredSumErrorMap()
	{
		return squaredErrorSumMap;
	}

	@Override
	public Map<String, Double> getRMSErrorMap()
	{
		return rmsErrorMap;
	}

	@Override
	public double getEpochRMSError()
	{
		return epochRMS;
	}

	@Override
	public void updateErrorSums(Map<String, Double> localErrorMap)
	{
		for (String name : localErrorMap.keySet())
		{
			double currentError = localErrorMap.get(name);
			double errorSum = totalErrorMap.containsKey(name) ? totalErrorMap.get(name) : 0;
			totalErrorMap.put(name, errorSum + currentError);

			double currentSquaredError = currentError * currentError;
			double squaredErrorSum = squaredErrorSumMap.containsKey(name) ? squaredErrorSumMap.get(name) : 0;
			squaredErrorSumMap.put(name, squaredErrorSum + currentSquaredError);
		}
	}

	@Override
	public void calculateEpochErrors(IDataSet dataSet)
	{
		for (String name : totalErrorMap.keySet())
		{
			double average = totalErrorMap.get(name) / dataSet.getDataSetSize();
			averageErrorMap.put(name, average);
		}

		for (String name : squaredErrorSumMap.keySet())
		{
			double rmsError = Math.sqrt(squaredErrorSumMap.get(name) / dataSet.getDataSetSize());
			rmsErrorMap.put(name, rmsError);
		}

		epochRMS = 0;
		for (Double value : squaredErrorSumMap.values())
		{
			epochRMS += value;
		}
		epochRMS /= (dataSet.getDataSetSize() * dataSet.getOutputNeuronNames().size());
		epochRMS = Math.sqrt(epochRMS);

	}

	@Override
	public void resetEpochExecutor()
	{
		epochIndex = 0;
		resetErrors();
	}

	@Override
	public void resetErrors()
	{
		totalErrorMap.clear();
		averageErrorMap.clear();
		squaredErrorSumMap.clear();
		rmsErrorMap.clear();

	}

}
